
# 6.863 -- Final Project
#
# Author: Capen Low, Andreea Bodnari
# 5/9/11
#
#


from train.ngramTrain import NgramTrain

import random

#
# The GridBuilder class is designed to take in a LanguageModel and size parameters (n,m)
#     to generate a grid of size n-by-m holding words in each cell, representing a
#     two-dimensional linguistic artifact.
#
# The given Probs.LanguageModel serves as an n-gram basis for populating the grid with
#     words in two dimensions.
# 

class GridBuilder:

    BASIC_MARKOV = "markov"
    COIN_FLIP = "coinflip"
    LIKELIER = "likelier"
    RARER = "rarer"


    METHOD = BASIC_MARKOV

    # initialize the GridBuilder construct
    def __init__(self, Method, HorizModel, VertModel, n, m):
        # Dimensions
        self.rows = n
        self.cols = m

        # Model
        self.HModel = HorizModel
        self.VModel = VertModel

        # Grid
        self.grid = Grid(n,m)

        # Blending Method
        self.method = Method
        
        # Output file
        self.gridOutput = "../experiments/Grid&Probs.txt"
        grid = open(self.gridOutput, "w")        
        grid.close()

    # Writes to file a given grid
    def writeGrid(self, aGrid, filename):
        f = open(filename, 'w')

        DISTANCE = 1
        
        grid = aGrid.grid

        for i in range(self.rows):
            for j in range(self.cols):
                word = grid[i][j][0]
                heur = grid[i][j][1]
                f.write(word + '(' + heur + ',' + DISTANCE + ')')

                if (j < self.cols - 1):
                    f.write(' ')
                else:
                    f.write('\n')

        f.close()

    # Prints a given grid
    def printGrid(self, aGrid):
        DISTANCE = 1

        grid = aGrid.grid

        for i in range(self.rows):
            rowstr = ''
            for j in range(self.cols):
                word = grid[i][j][0]
                heur = grid[i][j][1]
                rowstr = rowstr + word + '(' + str(heur) + ',' + str(DISTANCE) + ')'
                
                if (j < self.cols - 1):
                    rowstr = rowstr + ' '
                # print handles its own \n addition
                #else:
                    #rowstr = rowstr + '\n'
            print rowstr
     
    # Write to file a readable grid
    def writeReadableGrid(self, aGrid, words = True, bestPath = []):
        """
        Return a pretty string representation of a nxm matrix.

        @author: Andreea Bodnari
        @param matrix: Matrix representation (list with n rows of m elements).
        """
        grid = aGrid.grid
        lengths = []
          
        for row in range(self.rows):
            for column in range(self.cols):
                word = str(grid[row][column][0])
                
                cl = len(word) 
                try:
                    ml = lengths[column]
                    if cl > ml:
                        lengths[column] = cl
                except IndexError:
                    lengths.append(cl)
        
        lengths = tuple(lengths)
        format_string = ""
        for length in lengths:
            format_string += "%-" + str(length+4) + "s "
        format_string += "\n"
        
        matrix_str = ""
        for row in range(self.rows):
            toAppend = []
            for el in range(self.cols):
                if words :
                    if (row, el) in bestPath:
                        toAppend.append("++" +  grid[row][el][0]  + "++")
                    else:      
                        toAppend.append(grid[row][el][0])
                else:
                    if (row, el) in bestPath:
#                        toAppend.append('\033[0;31m%s\033[m' %grid[row][el][1])
                        toAppend.append("++"  + grid[row][el][1] + "++")
                    else:      
                        toAppend.append(grid[row][el][1])
            matrix_str += format_string % tuple(toAppend)
        
        fileOut = open(self.gridOutput, "a")
        fileOut.write( matrix_str)
        fileOut.write("\n")
        fileOut.close()
                                                                                                                

    # Creates and returns a new instance of a grid given the constraints within the
    #    GridBuilder object
    def constructGrid(self):
        
        # Clone the existing grid to create a working copy
        newGrid = self.grid

        # Populate Grid
        # We do this by iterating through the grid from [0,0] to [n-1,m-1],
        #  generating a 4-tuple of context containing the two words above and before
        #  the target cell, pass it to generateWord() and let it handle blending
        for i in range(self.rows):
            for j in range(self.cols):
#                print "Processing cell (" + str(i) + "," + str(j) + ")"
                context = [self.grid.getWord(i,j-2),
                           self.grid.getWord(i,j-1),
                           self.grid.getWord(i-2,j),
                           self.grid.getWord(i-1,j)]
                newGrid.grid[i][j] = self.generateGridTuple(context, self.method)

        # Return newly populated grid
        return newGrid


    # Generates a tuple containing [Word, HeuristicValue] given the surrounding context
    #   Takes in blendMethod flag that dictates type of blending mechanism used
    
    def generateGridTuple(self, context, blendMethod):

        # First cell is a special case: It just takes a new context-less word from HModel
        if (context == [None, None, None, None]):
#            print "Seeding..."
            newWord = self.HModel.gridSeed()
#            print "Seed word is " + newWord
            return [newWord, self.HModel.sequenceProbability([context[0], context[1], newWord])]

        # First row is a special case: Only generates on HModel
        elif (context[2:4] == [None, None]):
#            print "First row..."
            newHWord = self.HModel.wordGeneration([context[0],context[1]])
#            print "next word is " + newHWord
            return [newHWord, self.HModel.sequenceProbability([context[0],context[1],newHWord])]

        # First column is a special case: Only generates on VModel
        elif (context[0:2] == [None, None]):
#            print "First column..."
            newVWord = self.VModel.wordGeneration([context[2],context[3]])
#            print "next word is " + newVWord
            return [newVWord, self.VModel.sequenceProbability([context[2],context[3],newVWord])]


        # Else, use the blending method to create a new word
        else:
        
            # Basic Markov
            if (blendMethod == self.BASIC_MARKOV):
                return NgramTrain.wordGenerationMarkovModel(self.HModel, self.VModel, [context[0], context[1]], [context[2], context[3]])

            # Coin-Flip
            elif (blendMethod == self.COIN_FLIP):
                horizontalWord = self.HModel.wordGeneration([context[0],context[1]])
                verticalWord = self.VModel.wordGeneration([context[2],context[3]])
                
                coinval = random.random()
                if (coinval < 0.5):
                    return [horizontalWord, self.HModel.sequenceProbability([context[0],context[1],horizontalWord])]
                else:
                    return [verticalWord, self.VModel.sequenceProbability([context[2],context[3],verticalWord])]

            # Likelier: Returns the word with greater likelihood of two generated
            elif (blendMethod == self.LIKELIER):
                horizontalWord = self.HModel.wordGeneration([context[0],context[1]])
                verticalWord = self.VModel.wordGeneration([context[2],context[3]])
                
                hprob = self.HModel.sequenceProbability([context[0],context[1],horizontalWord])
                vprob = self.VModel.sequenceProbability([context[2],context[3],verticalWord])
                if (hprob >= vprob):
                    return [horizontalWord, hprob]
                else:
                    return [verticalWord, vprob]

            # Rarer: Returns the word with lesser likelihood of two generaetd
            elif (blendMethod == self.RARER):
                horiztonalWord = self.HModel.wordGeneration([context[0],context[1]])
                verticalWord = self.VModel.wordGeneration([context[2],context[3]])

                hprob = self.HModel.sequenceProbability([context[0],context[1],horizontalWord])
                vprob = self.VModel.sequenceProbability([context[2],context[3],verticalWord])
                if (hprob <= vprob):
                    return [horizontalWord, hprob]
                else:
                    return [verticalWord, vprob]
                                                                                

            # default -- shouldn't happen...
            else:
                return [None, 0.0]

class Grid:
    def __init__(self, n, m):
        # Dimensions
        self.rows = n
        self.cols = m

        # Grid, initialized as n rows, m columns, w/ the "None" value in all cells
        # Note that range(n) = (0, 1, 2, ..., n-2, n-1)
        newgrid = [None]  # Just a dummy declaration

        for i in range(n):
            # Create list of m elements, each reading "None"
            newrow = [None]
            for j in range(1,m):
                newrow.append(None)

            # If this is the first row, define grid, else append newrow to existing grid
            if (i == 0):
                newgrid = [newrow]
            else:
                newgrid.append(newrow)

        # Set self.grid to the the two-dimensional list of size n-by-m,
        #     populated with the value of "None" in all cells
        self.grid = newgrid


    # Fetches word stored in the grid.  Returns None if invalid coordinates are given.
    def getWord(self, i, j):
        if ((i < 0) or (j < 0)):
            return None
        elif ((i >= self.rows) or (j >= self.cols)):
            return None
        else:
            # As we are storing tuples in this grid, the word itself is stored in the
            #  first slot of the tuple.  getWord() returns this.
            return self.grid[i][j][0]
